#  -*- coding: utf-8 -*-

from pandas import DataFrame
import matplotlib.pyplot as plt
from data.data_module import DataModule
from trading.account import Account
from util.stock_util import get_trading_dates,get_code_name
from util.logger import QuantLogger
from util.database import DB_CONN
from factor.factor_module import FactorModule
import shutil
import numpy as np
from util.database import DB_CONN
from pymongo import UpdateOne
import os,time
from util.database import base_code_path

class Backtest:
    def __init__(self, strategy_option, begin_date=None, end_date=None):
        self.strategy_option = strategy_option
        if begin_date is None:
            self.begin_date = self.strategy_option.begin_date()
        else:
            self.begin_date = begin_date

        if end_date is None:
            self.end_date = self.strategy_option.end_date()
        else:
            self.end_date = end_date

        self.dm = DataModule()
        self.code_daily_cache = dict()

        self.logs_path = f"{base_code_path}/HunterQuant/logs"
        # if os.path.exists(logs_path):
        #     shutil.rmtree(logs_path)

        if False == os.path.exists(self.logs_path):
            os.mkdir(self.logs_path)


        self.logger = QuantLogger('backtest')
        self.candidatesLogger = QuantLogger('candidates')
        self.db = DB_CONN



    def start(self):
        """
        策略回测。结束后打印出收益曲线(沪深300基准)、年化收益、最大回撤
        """
        #策略名，来自策略文件的参数
        strategy_name = self.strategy_option.get_strategy_name()
        # 初始总资金
        initial_capital = self.strategy_option.capital()
        # 初始现金
        cash = initial_capital
        # 单只股票仓位上限
        single_position = self.strategy_option.single_position()

        # 从获取策略配置中获取股票池
        stock_pool = self.strategy_option.stock_pool()

        # 保存持仓股的日期
        account = Account()

        # 获取卖出信号
        sell_signal = self.strategy_option.sell_signal(account)

        # 获取买入信号
        buy_signal = self.strategy_option.buy_signal(account)

        # 时间为key的净值、收益和同期沪深基准
        df_profit = DataFrame(columns=['net_value', 'profit', 'hs300'])

        # 获取交易日历，
        #self.begin_date = "2017-04-13"
        #self.end_date = "2017-04-25"
        all_dates = get_trading_dates(begin_date=self.begin_date, end_date=self.end_date)

        # 获取沪深300在统计周期内的第一天的值
        hs300_k = self.dm.get_k_data('000300', index=True, begin_date=all_dates[0], end_date=all_dates[0])
        hs300_begin_value = hs300_k.loc[hs300_k.index[0]]['close']
        #获取上证指数的数据
        sz_k = self.dm.get_k_data('999999', index=True, begin_date=self.begin_date, end_date=self.end_date)
        sz_k.set_index(['date'], 1, inplace=True)

        # 获取股票池数据
        rebalance_dates, date_codes_dict = self.dm.get_strategy_option_data(strategy=strategy_name,begin_date=self.begin_date, end_date=self.end_date)
        #rebalance_dates = ["2017-04-14"]
        #date_codes_dict = {"2017-04-14":['600262']}

        fm = FactorModule()
        fm_daily = fm.get_single_stock_factors("999999", "hfq_ma", True, begin_date=self.begin_date, end_date=self.end_date)
        fm_daily.set_index(['date'], 1, inplace=True)
        # 获取止损策略
        stop_loss_policies = self.strategy_option.get_stop_loss(account)

        # 获取止盈策略
        stop_profit_policies = self.strategy_option.get_stop_profit(account)

        # 获取加仓策略
        add_position_policy = self.strategy_option.get_add_position(account)

        # 获取回测周期内股票池内所有股票的收盘价和前收价
        all_option_code_set = set()
        for rebalance_date in rebalance_dates:
            for code in date_codes_dict[rebalance_date]:
                all_option_code_set.add(code)

        # 缓存股票的日线数据
        for code in all_option_code_set:
            dailies_df = self.dm.get_k_data(code, index=False, autype=None, begin_date=self.begin_date, end_date=self.end_date)
            dailies_df.set_index(['date'], inplace=True)

            self.code_daily_cache[code] = dailies_df

        last_phase_codes = None
        this_phase_codes = None
        to_be_sold_codes = set()
        to_be_bought_codes = set()
        last_date = None

        # 加仓
        to_be_added_signals = dict()
        to_be_added_codes = set()

        suc_times = 0
        fail_times = 0
        total_suc_profit_ratio = 0
        total_fail_profit_ratio = 0
        total_rs = [0,0,0,0,0,0]

        # 按照日期一步步回测
        update_requests = []

        for _date in all_dates:
            self.logger.info('开始回测，日期：' + _date)
            fm_rs_date_df = fm.get_single_date_factors("rs",False,_date)
            fm_rs_date_df.set_index(['code'], 1, inplace=True)
            # 处理复权,更新持仓股的持仓数量
            account.adjust_holding_volume_at_open(last_date, _date)

            # 卖出
            if len(to_be_sold_codes) > 0:
                sold_codes_tmp = set(to_be_sold_codes)
                for code in sold_codes_tmp:
                    try:
                        if code in account.holding_codes:
                            holding_stock = account.get_holding(code)
                            holding_volume = holding_stock['volume']
                            sell_price = self.code_daily_cache[code].loc[_date]['open']

                            daily = self.code_daily_cache[code].loc[_date]
                            if False == daily["is_trading"]:
                                print('当日停牌，无法卖出，股票代码：%s, 日期： %s ' % (code, _date), flush=True)

                            else:
                                date_change_rate = self.db['change_rate'].find_one(
                                    {'code': code, 'index': False, 'date': _date},
                                    projection={"price_limit":True, '_id': False})
                                if date_change_rate is not None:
                                    price_limit = date_change_rate["price_limit"]

                                #不是跌停，-1跌停，-2 一字跌停，没有精细化考虑不是一字跌停时可以卖出的情况
                                if price_limit > -1:
                                    sell_amount = holding_volume * sell_price
                                    cash += sell_amount

                                    cost = holding_stock['cost']
                                    single_profit = round((sell_amount - cost) * 100 / cost,2)
                                    if single_profit > 0 :
                                        suc_times+=1
                                        total_suc_profit_ratio += single_profit
                                        # total_rs[0] += fm_rs_date_df.loc[code]["rs_60"]
                                        # total_rs[1] += fm_rs_date_df.loc[code]["rs_120"]
                                        # total_rs[2] += fm_rs_date_df.loc[code]["rs_250"]
                                    else:
                                        fail_times+=1
                                        total_fail_profit_ratio+=single_profit
                                        # total_rs[3] += fm_rs_date_df.loc[code]["rs_60"]
                                        # total_rs[4] += fm_rs_date_df.loc[code]["rs_120"]
                                        # total_rs[5] += fm_rs_date_df.loc[code]["rs_250"]

                                    update_requests.append(
                                        UpdateOne(
                                            {'code': code,
                                             'date': holding_stock['last_date']},
                                            {'$set': {
                                                    'code':code,
                                                    'date':holding_stock['last_date'],
                                                    'name': get_code_name(code,False),
                                                    'buy_date':holding_stock['buy_date'],
                                                    'volume':holding_stock['volume'],
                                                    'buy_price':holding_stock['buy_price'],
                                                    'sell_date':_date,
                                                    'sell_price':sell_price,
                                                    'profit':single_profit
                                                    }
                                            },
                                            upsert=True)
                                    )
                                    # 从持仓股中卖出
                                    account.sell_out(code)
                                    # 从代码列表中删除
                                    to_be_sold_codes.remove(code)

                                    print('卖出股票 %s,卖出数量： %6d,卖出价格： %6.2f,总额： %8.2f,单个股票收益： %4.2f' %
                                          (code, holding_volume, sell_price, sell_amount, single_profit))
                                    self.logger.info('卖出股票 %s,卖出数量： %6d,卖出价格： %6.2f,总额： %8.2f,单个股票收益： %4.2f',
                                          code, holding_volume, sell_price, sell_amount, single_profit)

                                else:
                                    print('当日跌停，无法卖出，股票代码：%s, 日期： %s，价格：%7.2f'
                                          % (code, _date, sell_price), flush=True)


                    except:
                        print('卖出时，发生异常：%s, %s' % (code, _date), flush=True)

            print('卖出后，现金: %10.2f' % cash)

            # 加仓逻辑
            add_codes_tmp = set(to_be_added_codes)
            for code in add_codes_tmp:
                add_signal = to_be_added_signals[code]
                try:
                    daily = self.code_daily_cache[code].loc[_date]
                    if False == daily["is_trading"]:
                        print('当日停牌，无法加仓，股票代码：%s, 日期： %s ' % (code, _date), flush=True)
                    else:
                        if cash > add_signal['position']:
                            buy_price = daily['open']

                            date_change_rate = self.db['change_rate'].find_one(
                                {'code': code, 'index': False, 'date': _date},
                                projection={"price_limit": True, '_id': False})
                            if date_change_rate is not None:
                                price_limit = date_change_rate["price_limit"]

                            #没有精细化考虑不是一字涨停可以买入的情况
                            if price_limit < 1:
                                volume = int(int(add_signal['position'] / buy_price) / 100) * 100
                                buy_amount = buy_price * volume
                                cash -= buy_amount
                                print('加仓股票 %s, 加仓数量：%6d,加仓价格： %6.2f,花费总额： %8.2f' % (code, volume, buy_price, buy_amount), flush=True)
                                self.logger.info('加仓股票 %s, 加仓数量：%6d,加仓价格： %6.2f,花费总额： %8.2f', code, volume, buy_price, buy_amount)
                                # 更新加仓后的持仓股
                                holding = account.get_holding(code)
                                holding['cost'] += buy_amount

                                #先用买入花费记录last_value,后续调用get_total_value时会更新为当日收盘时总额
                                holding['last_value'] += buy_amount

                                holding['volume'] += volume
                                holding['add_times'] += 1
                                holding['last_buy_hfq_price'] = buy_price * daily['au_factor']
                                account.update_holding(code, holding)

                                # 从待加仓列表中删除
                                to_be_added_codes.remove(code)
                                del to_be_added_signals[code]
                            else:
                                print('当日涨停，无法加仓，股票代码：%s, 日期： %s，价格：%7.2f'
                                      % (code, _date, buy_price), flush=True)

                except:
                    print('加仓时，发生错误：%s, %s' % (code, _date), flush=True)

            # 买入
            if len(to_be_bought_codes) > 0:
                sorted_to_be_bought_list = list(to_be_bought_codes)
                sorted_to_be_bought_list.sort()
                for code in sorted_to_be_bought_list:
                    try:
                        daily = self.code_daily_cache[code].loc[_date]
                        if False == daily["is_trading"]:
                            print('当日停牌，无法买入，股票代码：%s, 日期： %s ' % (code, _date), flush=True)
                        else:
                            if cash > single_position:
                                daily = self.code_daily_cache[code].loc[_date]
                                buy_price = daily['open']

                                date_change_rate = self.db['change_rate'].find_one(
                                    {'code': code, 'index': False, 'date': _date},
                                    projection={"price_limit": True, '_id': False})
                                if date_change_rate is not None:
                                    price_limit = date_change_rate["price_limit"]

                                if price_limit < 1:
                                    volume = int(int(single_position / buy_price) / 100) * 100
                                    buy_amount = buy_price * volume
                                    cash -= buy_amount
                                    print('买入股票 %s, 买入数量：%6d,买入价格： %6.2f,花费总额： %8.2f' % (code, volume, buy_price, buy_amount), flush=True)
                                    self.logger.info('买入股票 %s, 买入数量：%6d,买入价格： %6.2f,花费总额： %8.2f', code, volume, buy_price, buy_amount)

                                    # 维护账户的持仓股
                                    account.buy_in(code, buy_price=buy_price, volume=volume, cost=buy_amount,date=_date,last_date=last_date)

                                    # 如果加仓策略不为空，则更新持仓股
                                    if add_position_policy is not None:
                                        holding = account.get_holding(code)
                                        holding['last_buy_hfq_price'] = buy_price * daily['au_factor']
                                        add_position_policy.update_holding(code, _date, holding)
                                else:
                                    print('当日涨停，无法买入，股票代码：%s, 日期： %s，价格：%7.2f'
                                          % (code, _date, buy_price), flush=True)

                    except:
                        print('买入时，发生错误：%s, %s' % (code, _date), flush=True)

            print('买入后，现金: %10.2f' % cash)
            self.logger.info('买入后，现金: %10.2f', cash)

            # 计算总市值
            total_value = account.get_total_value(_date)
            # 计算总资产
            total_capital = total_value + cash

            print('收盘后，现金: %10.2f, 总资产: %10.2f' % (cash, total_capital))
            self.logger.info('收盘后，现金: %10.2f, 总资产: %10.2f', cash, total_capital)

            # 持仓股代码列表
            holding_codes = account.holding_codes
            # 如果调整日，则获取新一期的股票列表
            if _date in rebalance_dates:
                # 暂存为上期的日期
                if this_phase_codes is not None:
                    last_phase_codes = this_phase_codes
                this_phase_codes = date_codes_dict[_date]

            # 检查是否有需要第二天卖出的股票
            if sell_signal is not None:
                for holding_code in holding_codes:
                    if sell_signal.is_match(holding_code, _date):
                        to_be_sold_codes.add(holding_code)
                        print('发出卖出信号，股票：%s' % holding_code, flush=True)


            # 检测止损信号
            if stop_loss_policies is not None:
                for holding_code in holding_codes:
                    for stop_loss_policy in stop_loss_policies:
                        if stop_loss_policy.is_stop(holding_code, _date):
                            to_be_sold_codes.add(holding_code)
                            print('止损，股票：%s' % holding_code, flush=True)
                        else:
                            stop_loss_policy.update_holding(holding_code, _date)

            # 检测止盈信号
            if stop_profit_policies is not None:
                for holding_code in holding_codes:
                    for stop_profit_policy in stop_profit_policies:
                        if stop_profit_policy.is_stop(holding_code, _date):
                            to_be_sold_codes.add(holding_code)
                            print('止盈，股票：%s' % holding_code, flush=True)
                        else:
                            stop_profit_policy.update_holding(holding_code, _date)

            self.logger.info('待卖股票，日期：%s，代码列表：%s', _date, to_be_sold_codes )

            # 检测是否有需要加仓的股票
            if add_position_policy is not None:
                for holding_code in holding_codes:
                    add_signal = add_position_policy.get_add_signal(holding_code, _date)
                    if add_signal is not None:
                        to_be_added_signals[holding_code] = add_signal
                        to_be_added_codes.add(holding_code)

            # 检查是否有需要第二天买入的股票,没有买入信号则买入本期所有股票，有买入信号按买入信号执行
            to_be_bought_codes.clear()
            if this_phase_codes is not None:
                for _code in this_phase_codes:
                    if _code not in holding_codes:

                    # if _code not in holding_codes and \
                    #          _code in fm_rs_date_df.index and \
                    #         fm_rs_date_df.loc[_code]["rs_60"] > 30 and \
                    #         fm_rs_date_df.loc[_code]["rs_120"] > 30 and \
                    #         fm_rs_date_df.loc[_code]["rs_250"] > 30 :

                        if buy_signal is not None :
                            if buy_signal.is_match(_code, _date):
                                to_be_bought_codes.add(_code)
                        else:
                            to_be_bought_codes.add(_code)


            self.logger.info('待买股票，日期：%s，代码列表：%s', _date, to_be_bought_codes)
            self.candidatesLogger.info('待买股票，日期：%s，代码列表：%s', _date, to_be_bought_codes)



            # 计算沪深300的增长
            hs300_k_current = self.dm.get_k_data('000300', index=True, begin_date=_date, end_date=_date)
            hs300_current_value = hs300_k_current.loc[hs300_k_current.index[0]]['close']

            last_date = _date
            df_profit.loc[_date] = {
                'net_value': round(total_capital / initial_capital, 2),
                'profit': round(100 * (total_capital - initial_capital) / initial_capital, 2),
                'hs300': round(100 * (hs300_current_value - hs300_begin_value) / hs300_begin_value, 2)
            }

        #输出回测结果

        # 批量写入，提高访问效率
        collection = DB_CONN["%s_backtest"%strategy_name]
        collection.create_index([('code', 1), ('date', 1)])
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存回测结果到数据集：%s，插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, update_result.upserted_count, update_result.modified_count,
                   (end_time - start_time)),
                  flush=True)

        # 打印回测收益曲线数值
        print('Profit history start')
        for index_date in df_profit.index:
            print('回测日期：%s,策略收益： %6.2f,沪深300收益： %6.2f' %
                  (index_date, df_profit.loc[index_date]['profit'], df_profit.loc[index_date]['hs300']),
                  flush=True)
        print('Profit history end')

        drawdown = self.compute_drawdown(df_profit['net_value'])
        annual_profit, sharpe_ratio = self.compute_sharpe_ratio(df_profit['net_value'])

        if suc_times != 0:
            avg_suc_profit_ratio = total_suc_profit_ratio/suc_times
        else:
            avg_suc_profit_ratio = 0

        if fail_times != 0:
            avg_fail_profit_ratio = total_fail_profit_ratio/fail_times
        else:
            avg_fail_profit_ratio = 0

        suc_ratio = suc_times/(suc_times+fail_times) * 100
        if avg_fail_profit_ratio != 0:
            risk_to_profit_ratio = abs(avg_suc_profit_ratio/avg_fail_profit_ratio)
        else:
            #无风险
            risk_to_profit_ratio = 1E8

        print('回测结果 %s - %s，年化收益： %7.3f, 最大回撤：%7.3f, 夏普比率：%4.2f,交易次数：%d, 成功次数：%d,失败次数: %d, 成功率：%.2f,平均收益：%.2f,平均损失：%.2f,风险收益比：%.2f' %
              (self.begin_date, self.end_date, annual_profit, drawdown, sharpe_ratio,(suc_times+fail_times), suc_times,fail_times,suc_ratio,avg_suc_profit_ratio,avg_fail_profit_ratio,risk_to_profit_ratio))
        #print("RS:60:%d,120:%d,250:%d,60:%d,120:%d,250:%d" % (total_rs[0],total_rs[1],total_rs[2],total_rs[3],total_rs[4],total_rs[5]))
        df_profit.plot(title='Backtest Result', y=['profit', 'hs300'], kind='line')

        # plt.figure(figsize=(5,5))
        # ax = plt.subplot(111)
        # ax.set_xlim(self.begin_date,self.end_date)
        # ax.plot(df_profit['profit'], linestyle='--', marker='o', markersize=5)
        # ax.plot(df_profit['hs300'], linestyle='-')
        # ax.set_title(u'回测结果')
        # ax.set_ylabel(u'收益率')
        plt.show()

    def compute_sharpe_ratio(self, net_values):
        """
        计算夏普比率
        :param net_values: 净值列表
        """

        # 总交易日数
        trading_days = len(net_values)
        # 所有收益的DataFrame
        profit_df = DataFrame(columns={'profit'})
        # 收益之后，初始化为第一天的收益
        profit_df.loc[0] = {'profit': round((net_values[0] - 1) * 100, 2)}
        # 计算每天的收益
        for index in range(1, trading_days):
            # 计算每日的收益变化
            profit = (net_values[index] - net_values[index - 1]) / net_values[index - 1]
            profit = round(profit * 100, 2)
            profit_df.loc[index] = {'profit': profit}

        # 计算标准差
        profit_std = pow(profit_df.var()['profit'], 1 / 2)

        # 年化收益
        annual_profit = self.compute_annual_profit(trading_days, net_values[-1])

        # 夏普比率
        sharpe_ratio = (annual_profit - 4.75) / profit_std

        return annual_profit, sharpe_ratio

    def compute_drawdown(self, net_values):
        """
        计算最大回撤
        :param net_values: 净值列表
        """
        # 最大回撤初始值设为0
        max_drawdown = 0
        index = 0
        # 双层循环找出最大回撤
        for net_value in net_values:
            #从index向后查询最大回撤，然后逐步index+1，进行更新
            for sub_net_value in net_values[index:]:
                drawdown = 1 - sub_net_value / net_value
                if drawdown > max_drawdown:
                    max_drawdown = drawdown

            index += 1

        return max_drawdown*100

    def compute_annual_profit(self, trading_days, net_value):
        """
        计算年化收益
        """

        annual_profit = 0
        if trading_days > 0:
            # 计算年数
            years = trading_days / 245
            # 计算年化收益
            annual_profit = pow(net_value, 1 / years) - 1

        annual_profit = round(annual_profit * 100, 2)

        return annual_profit

    def find_out_stocks(self, last_phase_codes, this_phase_codes):
        """
        找到上期入选本期被调出的股票，这些股票将必须卖出
        :param last_phase_codes: 上期的股票列表
        :param this_phase_codes: 本期的股票列表
        :return: 被调出的股票列表
        """
        out_stocks = []

        for code in last_phase_codes:
            if code not in this_phase_codes:
                out_stocks.append(code)

        return out_stocks

